package com.hexiang.ratelimit.aspect;

import com.google.common.util.concurrent.RateLimiter;
import com.hexiang.ratelimit.RateLimitException;
import com.hexiang.ratelimit.annotation.RateLimit;
import com.hexiang.ratelimit.store.RateLimiterStore;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.lang.reflect.Method;

@Aspect
@Slf4j
@Order(0)
@Component
public class RateLimitAspect {

    @Resource
    private RateLimiterStore rateLimiterStore;

    @Around("@annotation(com.hexiang.ratelimit.annotation.RateLimit)")
    public Object invokeResourceWithRateLimit(ProceedingJoinPoint pjp) throws Throwable {
        Method originMethod = resolveMethod(pjp);

        RateLimit annotation = originMethod.getAnnotation(RateLimit.class);
        if (annotation == null) {
            // Should not go through here.
            throw new IllegalStateException("Wrong state for RateLimit annotation");
        }
        Class<?> targetClass = pjp.getTarget().getClass();

        String apiKey = targetClass.getName() + ":" + originMethod.getName();

        RateLimiter rateLimiter = rateLimiterStore.resolveRateLimit(apiKey,annotation.permitsPerSecond());

        if (rateLimiter.tryAcquire()) {
            return pjp.proceed();
        } else {
            throw new RateLimitException(apiKey, rateLimiter.getRate());
        }
    }

    Method resolveMethod(ProceedingJoinPoint joinPoint) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Class<?> targetClass = joinPoint.getTarget().getClass();

        Method method = getDeclaredMethodFor(targetClass, signature.getName(),
                signature.getMethod().getParameterTypes());
        if (method == null) {
            throw new IllegalStateException("Cannot resolve target method: " + signature.getMethod().getName());
        }
        return method;
    }

    private Method getDeclaredMethodFor(Class<?> clazz, String name, Class<?>... parameterTypes) {
        try {
            return clazz.getDeclaredMethod(name, parameterTypes);
        } catch (NoSuchMethodException e) {
            Class<?> superClass = clazz.getSuperclass();
            if (superClass != null) {
                return getDeclaredMethodFor(superClass, name, parameterTypes);
            }
        }
        return null;
    }

}